import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import numpy as np
import cv2 as cv
import torchvision.transforms as transforms

class GradCAM(object):
    '''gradCAM '''
    
    def __init__(self,model_dict,verbose=False):
        self.model_type=model_dict["type"] # 
        self.target_layer=model_dict["target_layer"] # 获得图层 A 
        self.model=model_dict["model"] # 模型 
        self.gradients=dict()
        self.activations=dict()

        # 声明 函数，用于在前向传播中，捕捉输出，后向传播中，捕捉梯度
        def forward_hook(module,input_act,output):
            self.activations["input"]=input_act
            self.activations["value"]=output
            return None
        def backward_hook(module,grad_in,grad_out):
            self.gradients["input"]=grad_in
            self.gradients["value"]=grad_out

        # 注册 函数
        self.target_layer.register_forward_hook(forward_hook)
        self.target_layer.register_backward_hook(backward_hook)
    
    def calGradhot(self,input,retain_graph=True,clsn=2): 
        """计算 gradCAM """
        """
        args:
            input:输入的图片
            class_index: 需要处理的类别
        """
        result=[]
        b,c,h,w=input.size() # 获取图片的batch，channels,height,width
        imgout=self.model(input) # 完成的图片的前向传播计算，获取 特征图像 A 的值
        for index in range(b):
            tempinfo={"b":index,'label':imgout[index,:],"saliency":[]}
            for cidx in range(clsn):
                score=imgout[index,cidx].squeeze() # Yc
                self.model.zero_grad() # 清除模型中的 梯度计算值
                score.backward(retain_graph=True) # 后向传播计算，并保留计算结果，获取导数 dY/dA
                activate=self.activations["value"]
                #print(self.gradients["value"].shape)
                gradient=self.gradients["value"][0]# 注意这里为0
                b,k,u,v=gradient.size()
                
                #  计算 a
                alpha=gradient.view(b,k,-1).mean(2).view(b,k,1,1)


                saliency_map=(alpha*activate).sum(1,keepdim=True) # 计算 Lc
                saliency_map=F.relu(saliency_map)
                saliency_map=F.upsample(saliency_map,size=(h,w),mode="bilinear",align_corners=False)
                # 计算归一化
                saliency_map_min, saliency_map_max = saliency_map[index,:,:,:].min(), saliency_map[index,:,:,:].max()
                tempinfo["saliency"].append({'data':saliency_map[index,:,:,:],'min':saliency_map_min,"max":saliency_map_max})
                #saliency_map = (saliency_map - saliency_map_min).div(saliency_map_max - saliency_map_min).data    

            # 计算两者之间的归一化
            # 确定最小值
            minvalue=tempinfo["saliency"][0]["min"]
            maxvalue=tempinfo["saliency"][0]["max"]
            for temps in tempinfo["saliency"]:
                if temps["min"]<minvalue:
                    minvalue=temps["min"]
                if temps["max"]>maxvalue:
                    maxvalue=temps["max"]  
            # 计算归一化
            for i in range(len(tempinfo["saliency"])):
                tempinfo["saliency"][i]["data"]=(tempinfo["saliency"][i]["data"]-minvalue).div(maxvalue-minvalue).data
                
            result.append(tempinfo)    
        return result

class GradCAMpp(GradCAM):
    def __init__(self,model_dict,verbose=False):
        super(GradCAMpp,self).__init__(model_dict,verbose=verbose)
    
    def calGradhot(self,input,retain_graph=True,clsn=3):

        result=[] # 计算结果
        b,c,h,w=input.size() # 批次，通道，高，宽
        imgout=self.model(input) # 计算输出
        # 逐批次计算归纳图
        for bi in range(b):
            tempinfo={"b":bi,'label':imgout[bi,:],"saliency":[]}
            for cls_ids in range(clsn):
                score=imgout[bi,cls_ids] # 获得类别
                self.model.zero_grad() # 反向求导归零---注意这个方法，只是适用于 为线性或者 relu 激活函数
                score.backward(retain_graph=retain_graph) # 反向传播
                activate=self.activations["value"] # 获得激活函数
                gradient=self.gradients["value"][0]# 注意这里为0 获取 求解结果 A
                b,k,u,v=gradient.size() # 导数层的求导结果 dS/dA

                alpha_num=gradient.pow(2) # 分子
                alpha_denum=gradient.pow(2).mul(2)+activate.mul(gradient.pow(3)).view(b,k,u*v).sum(-1,keepdim=True).view(b,k,1,1) # 分母
                alpha_denum=torch.where(alpha_denum!=0.0,alpha_denum,torch.ones_like(alpha_denum)) # 保证全部用1 填充，防止除0
                alpha=alpha_num.div(alpha_denum+1e-7) # +1e-7 保证结尾不为0

                positive_gradient=F.relu(score.exp()*gradient) # ReLU(dY/dA) =ReLU(exp(s)/dS/dA)
                
                # 计算权重
                weight=(alpha*positive_gradient).view(b,k,u*v).sum(-1).view(b,k,1,1) # W

                saliency_map=(weight*activate).sum(1,keepdim=True) # 计算 Lc
                saliency_map=F.relu(saliency_map)
                saliency_map=F.upsample(saliency_map,size=(h,w),mode="bilinear",align_corners=False)

                # 计算归一化
                saliency_map_min, saliency_map_max = saliency_map[bi,:,:,:].min(), saliency_map[bi,:,:,:].max()
                #saliency_map = (saliency_map-saliency_map_min).div(saliency_map_max-saliency_map_min).data

                tempinfo["saliency"].append({'data':saliency_map[bi,:,:,:],'min':saliency_map_min,"max":saliency_map_max})
            # 计算两者之间的归一化
            # 确定最小值
            minvalue=tempinfo["saliency"][0]["min"]
            maxvalue=tempinfo["saliency"][0]["max"]
            for temps in tempinfo["saliency"]:
                if temps["min"]<minvalue:
                    minvalue=temps["min"]
                if temps["max"]>maxvalue:
                    maxvalue=temps["max"]       

            # 计算归一化i
            for i in range(len(tempinfo["saliency"])):
                tempinfo["saliency"][i]["data"]=(tempinfo["saliency"][i]["data"]-minvalue).div(maxvalue-minvalue).data

            result.append(tempinfo)  
        return result

        